import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split

# 假设data是已加载的DataFrame
data = pd.read_excel("F:\\code\\data\\ask\\long_tail_data_least5.xlsx",usecols=[ 'f_id','doctor_id','add_time'])

# 1. 删除交互医生个数小于5的用户的所有交互记录
interaction_counts = data.groupby('f_id')['doctor_id'].nunique()
valid_users = interaction_counts[interaction_counts >= 5].index
filtered_data = data[data['f_id'].isin(valid_users)]

# 2. 重新编码f_id和doctor_id并保存对应关系
f_id_mapping = {old_id: new_id for new_id, old_id in enumerate(filtered_data['f_id'].unique())}
doctor_id_mapping = {old_id: new_id for new_id, old_id in enumerate(filtered_data['doctor_id'].unique())}

filtered_data['f_id'] = filtered_data['f_id'].map(f_id_mapping)
filtered_data['doctor_id'] = filtered_data['doctor_id'].map(doctor_id_mapping)

f_id_df = pd.DataFrame(list(f_id_mapping.items()), columns=['old_f_id', 'new_f_id'])
doctor_id_df = pd.DataFrame(list(doctor_id_mapping.items()), columns=['old_doctor_id', 'new_doctor_id'])

# 保存映射关系
f_id_df.to_csv('f_id_mapping.csv', index=False)
doctor_id_df.to_csv('doctor_id_mapping.csv', index=False)

# 3. 按照7:3的比例划分训练集和测试集
train_file = open('train.txt', 'w')
test_file = open('test.txt', 'w')

for user_id, user_data in filtered_data.groupby('f_id'):
    sorted_data = user_data.sort_values('add_time')
    unique_doctors = sorted_data['doctor_id'].unique()

    train_size = int(len(unique_doctors) * 0.8)
    train_doctors = unique_doctors[:train_size]
    test_doctors = unique_doctors[train_size:]

    train_file.write(f"{user_id} " + " ".join(map(str, train_doctors)) + "\n")
    test_file.write(f"{user_id} " + " ".join(map(str, test_doctors)) + "\n")

train_file.close()
test_file.close()
